'''
1. 预处理-选择base，诊断病灶
2. 特征点检测
3. 特征描述
4. 与基底图特征点匹配
5. 去除错误匹配点
6. 计算与基底图的转移矩阵
7 交叉验证
8. 可视化
'''
import os
import json
from PIL import Image
import numpy as np
from configs import get_configs
from util.lesion_detector import ROP_single_clinic
from util.visualization_multi import Visualizer
visualize_type = '2D'  # 2D or 3D or None

visual_key_point = True
visual_match = True


def base_selector(optic_disc_coor_list, image_resolu):
    '''
    选择optic_disc 最接近中心位置的为base

    optic_disc_coor_lisst (list): 一个列表,每个元素是 {
            "position": [ ，],  # 中心位置
            "distance": "visible" | "near" | "far"  # 距离
        },
    '''
    center_x = image_resolu[0]//2
    center_y = image_resolu[1]//2
    res = -1
    min_dis = float('inf')
    for i, optic_disc in enumerate(optic_disc_coor_list):
        if optic_disc['distance'] == 'visible':
            # 计算和center的距离
            x = optic_disc['position'][0]
            y = optic_disc['position'][1]
            distance = (center_x-x)**2+(center_y-y)**2
            if distance < min_dis:
                min_dis = distance
                res = i
    if res == -1:
        raise ValueError("No visible optic disc found")
    return res


if __name__ == '__main__':
    # 获取参数
    args = get_configs()
    if visualize_type:
        os.makedirs(os.path.join(args.experiments_dir,
                    visualize_type), exist_ok=True)
    # 读取文件
    with open(os.path.join(args.data_path, 'annotations.json')) as f:
        data_dict = json.load(f)
    # with open(os.path.join(args.data_path,'follo_up.json')) as f:
    #     follow_up_dict=json.load(f)
    # with open(os.path.join(args.data_path,'split',f'{args.split_name}.json')) as f:
    #     split_dict=json.load(f)
    
    # debug mode
    follow_up_dict = {
        "tmp": [
            {
                "result": {
                    "OS": [
                    "1243.jpg",
                    "1244.jpg",
                    "1245.jpg",
                    "1246.jpg",
                    "1247.jpg"
                ]
                }
            },
        ]

    }
    split_dict = {'test': ['tmp']}

    
    if visual_key_point or visual_match:
        os.makedirs(os.path.join(args.experiments_dir,
                    'visual_detail'), exist_ok=True)
        
    # 实例化模型
    lesion_detecor = ROP_single_clinic(
        image_resolution=(args.images_resolu, args.images_resolu),
        threshold=args.lesion_threshold,
        bbox_expand=args.bbox_expand
    )

    # 实例化特征匹配器
    feature_matcher = FeatureMatcher(
        match_feature_num=2000,
        scaleFactor=1.2,
        nlevels=8
    )

    # 实例化交叉验证器
    cross_validator = LesionCrossValidator(
        base_image_shape=(args.images_resolu, args.images_resolu)
    )
    os.makedirs(os.path.join(args.experiments_dir, 'keypoints'), exist_ok=True)
    # 按照hid进行诊断
    for hid in split_dict['test']:
        for record_rank, follow_up_record in enumerate(follow_up_dict[hid]):
            for eye in ['OS', 'OD']:
                if eye in follow_up_record["result"]:

                    # 获取单张眼睛的图像列表
                    image_name_list = follow_up_record["result"][eye]
                    if len(image_name_list) < 2:
                        continue  # 图片数量太少，无法进行交叉验证
                    if visual_key_point or visual_match:
                        visual_detail_dir = os.path.join(
                            args.experiments_dir, 'visual_detail', f'{hid}_{str(record_rank+1)}_{eye}')
                        os.makedirs(visual_detail_dir, exist_ok=True)

                    # 选择base
                    base_idx = base_selector(
                        optic_disc_coor_list=[
                            data_dict[image_name]['optic_disc_pred'] for image_name in image_name_list],
                        image_resolu=args.dataset_info.image_resolution)
                    base_image_path = os.path.join(
                        args.data_path, data_dict[image_name_list[base_idx]]['image_path'])
                    feature_matcher._set_base_image(
                        base_image_path=os.path.join(
                            args.data_path, data_dict[image_name_list[base_idx]]['vascular_path'])
                    )

                    # 开始交叉验证
                    lesion_coords_list = []
                    transform_matrices = []
                    detect_success_list = []  # 检测的特征点过少，证明无法匹配
                    lesion_discover = False
                    for i, image_name in enumerate(image_name_list):
                        if i == base_idx:
                            continue
                        data = data_dict[image_name]
                        if "ridge_seg" not in data:
                            raise ValueError(
                                f"Ridge segmentation result not found in {image_name}, please use ridge_segmentation repository to generate ridge  segmentation masks")
                        if data['ridge_seg']['max_val'] > 0.5:
                            ridge_seg_path = os.path.join(
                                args.data_path, data['ridge_seg']['ridge_seg_path'])
                            if not os.path.exists(ridge_seg_path):
                                raise ValueError(
                                    f"Ridge segmentation file not found: {ridge_seg_path} from {image_name}, please use ridge_segmentation repository to generate ridge  segmentation masks")
                            lesion_bboxes = lesion_detecor.get_lesion_location(
                                ridge_seg_path)
                            lesion_discover = True
                        else:
                            lesion_bboxes = []
                        # 读取图像
                        image = Image.open(os.path.join(args.data_path, data['vascular_path'])).resize(
                            (args.images_resolu, args.images_resolu))
                        image = np.array(image)

                        # 检测特征点
                        keypoints, descriptors = feature_matcher.detect_features(
                            image)
                        if visual_key_point:
                            feature_matcher.visual_keypoints(
                                image_path=os.path.join(
                                    args.data_path, data['vascular_path']),
                                keypoints=keypoints,
                                save_path=os.path.join(
                                    visual_detail_dir, image_name)
                            )
                            # raise ValueError("Stop")
                        # 匹配特征点
                        if len(keypoints) < args.key_point_threshold:
                            continue
                        detect_success_list.append(image_name)
                        matches = feature_matcher.match_features(descriptors)
                        # 去除错误匹配点
                        good_matches = feature_matcher.filter_matches(matches)

                        if visual_match:
                            feature_matcher.visual_matches(
                                image_path=os.path.join(
                                    args.data_path, data['vascular_path']),
                                keypoints=keypoints,
                                good_matches=good_matches,
                                save_path=os.path.join(
                                    visual_detail_dir, image_name)
                            )
                            # raise ValueError("Stop")
                        # 计算转移矩阵
                        H = feature_matcher.calculate_transform_matrix(
                            keypoints, good_matches)

                        lesion_coords_list.append(lesion_bboxes)
                        transform_matrices.append(H)
                    if not lesion_discover:
                        continue
                    # 交叉验证
                    validate_lesions, mapped_image_polygons, mapped_lesion_polygons = cross_validator.get_validated_lesions(
                        lesion_coords_list=lesion_coords_list,
                        transform_matrices=transform_matrices
                    )
                    # 可视化
                    if visualize_type == '2D':
                        visualizer = Visualizer(
                            base_image_path=base_image_path,
                            validated_lesions=validate_lesions,
                            transform_matrices=transform_matrices,
                            image_paths_list=[os.path.join(
                                args.data_path, data_dict[image_name]['image_path']) for image_name in detect_success_list],
                            mapped_image_polygons=mapped_image_polygons,
                            image_resolution=(
                                args.images_resolu, args.images_resolu)
                        )
                        visualizer.visualize_2D(save_path=os.path.join(
                            args.experiments_dir, visualize_type, f'{hid}_{str(record_rank+1)}_{eye}.obj'))
                    elif visualize_type == '3D':
                        visualizer = Visualizer(
                            base_image_path=base_image_path,
                            validated_lesions=validate_lesions,
                            transform_matrices=transform_matrices,
                            image_paths_list=[os.path.join(
                                args.data_path, data_dict[image_name]['image_path']) for image_name in detect_success_list],
                            mapped_image_polygons=mapped_image_polygons,
                            image_resolution=(
                                args.images_resolu, args.images_resolu)
                        )
                        visualizer.visualize_3D(save_path=os.path.join(
                            args.experiments_dir, visualize_type, f'{hid}_{str(record_rank+1)}_{eye}.obj'))
                    else:
                        pass
